# 实战,训练一个分类器
"""
对于图像，Pillow，OpenCV 等包很有用
对于音频，请使用 SciPy 和 librosa 等包
对于文本，基于 Python 或 Cython 的原始加载，或者 NLTK 和 SpaCy 很有用
"""
# %%
# 先设置好路径
import pathlib
from pathlib import Path

# 当前文件的目录路径
CUR_FILE_DIR = Path(__file__).parent
# 强制绝对路径，在shell中执行时才用，如果执行文件这条应该注释
# 这里需要注意，因为是在交互式运行的
# 所有无法通过Path(__file__).parent获取到当前的路径
# CUR_FILE_DIR = Path(r'C:\files\git仓库\pytorch-learning\pytorch学习\快速入门')

# %%
""" 加载并标准化CIFAR10 """
import torch
import torchvision
import torchvision.transforms as transforms

# torchvision的数据集是[0,1]范围的PILImage
# 要归一化到[-1,1]范围
# compose，就是将几个transfroms集合到1个中
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 下载训练数据集,设置存放的root目录,然后进行数据转换
# 下载训练数据集，要设置train = True
# 把路径设为当前目录路径下的/data，并转换成str
trainset = torchvision.datasets.CIFAR10(root=str(CUR_FILE_DIR.joinpath('data')),
                                        train=True,
                                        download=True,
                                        transform=transform)
# 数据加载器,喂入数据集,批量大小,训练集肯定要打乱，测试集不用打乱
# 设置读取数据的线程数量，如果在windows上出现bug那么设置成0就可以解决问题
trainloader = torch.utils.data.DataLoader(trainset,
                                          batch_size=4,
                                          shuffle=True,
                                          num_workers=0)
# 下载测试数据集,设置存放的root目录,然后进行数据转换
testset = torchvision.datasets.CIFAR10(root=str(CUR_FILE_DIR.joinpath('data')),
                                       train=False,
                                       download=True,
                                       transform=transform)
# 数据加载器,设置数据集,批量大小,是否打乱,
# 一般来说测试集是不打乱的:为什么呢?为了复现效果
# 设置读取数据的线程数量，如果在windows上出现bug那么设置成0就可以解决问题
testloader = torch.utils.data.DataLoader(testset,
                                         batch_size=4,
                                         shuffle=False,
                                         num_workers=0)
"""定义标签"""
classes = ('plane', 'car', 'bird', 'cat', 'deer',
           'dog', 'frog', 'horse', 'ship', 'truck')
# %%
# 来展示一些训练图像吧
import matplotlib.pyplot as plt
import numpy as np


# 定义展示图像
def imshow(img):
    # 这个img是tensor
    # 反向去掉归一化
    img = img / 2 + 0.5
    # 把tensor转换成numpy
    npimg = img.numpy()
    # 展示,但要先进行转置
    # 原本的0，1，2现在变成1，2，0 为什么？
    # 因为tensor是chw，而真正的图片是hwc
    # 其实就是pytorch的样本表示方式为(n,c,h,w)把通道数放在前面了
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


# 取得一些随机的训练图片
# 先把数据载入器弄成迭代器
dataiter = iter(trainloader)
# 用next从迭代器中取出第1个样本
images, labels = next(dataiter)
# 展示图片,这里的make_grid是干啥的？
# make_grid顾名思义制造网格，把几张照片制作成一个网格
# make_grid的作用是将若干幅图像拼成一幅图像
imshow(torchvision.utils.make_grid(images))
# 打印标签
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))
# %%
# 定义卷积神经网络
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# 实例化网络
net = Net()
# 看看网络长什么样子
print(net)
# %%
# 定义损失函数和优化器
import torch.optim as optim

# 使用分类交叉熵
criterion = nn.CrossEntropyLoss()
# 定义优化器带动量的SGD
# optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
"""Accuracy:53%"""
# 第2次我会改成Adam会不会好一些呢
# optimizer = optim.Adam(net.parameters())
"""Accuracy:55%"""
# 第3次我会改成RmsProp会不会好一些呢
optimizer = optim.Adam(net.parameters())
"""Accuracy:55%"""
"""没有太大的区别，RmsProp和Adam都足够好！"""
# %%
# 训练网络
# 第一层肯定是epoch
for epoch in range(2):
    # 2个epoch，遍历数据集两遍
    running_loss = 0.0
    # enumerate(iterable,0)是什么意思呢？
    for i, data in enumerate(trainloader, 0):
        # 取得输入，数据是(输入，标签）的元组，
        # 先解包，而且是1个batch的数据
        inputs, labels = data
        # 将梯度置零
        optimizer.zero_grad()
        # 前向传播+反向传播+更新权重
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        # 打印统计数据
        running_loss += loss.item()
        if i % 2000 == 1999:
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0
print('Finished Training')
# %%
# 保存训练好的模型
PATH = str(CUR_FILE_DIR.joinpath('cifar_net.pth'))
# .pth文件是什么？
# pytorch的模型文件,save(model.state_ditc(),PATH)是只保存参数
# 其实dict就在提醒，这是一个参数字典
torch.save(net.state_dict(), PATH)

# %%
# 测试网络
# 取1个样本测试一下看看呢
dataiter = iter(testloader)
images, labels = next(dataiter)
# 打印出图片和标签
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))
# %%
# 加载模型
net = Net()
net.load_state_dict(torch.load(PATH))
# 进行预测
outputs = net(images)
# _表示占位没有实际意义
_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
                              for j in range(4)))

# %%
# 看看在整个数据集上的表现
correct = 0
total = 0
# 表示被这个warp起来的语句，不需要构建计算图，用于预测
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
        100 * correct / total))
# Accuracy of the network on the 10000 test images: 55 %的精度
# %%
# 哪些类的表现良好，哪些类的表现不佳：
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predicted = torch.max(outputs, 1)
        c = (predicted == labels).squeeze()
        for i in range(4):
            label = labels[i]
            class_correct[label] += c[i].item()
            class_total[label] += 1

for i in range(10):
    print('Accuracy of %5s : %2d %%' % (
        classes[i], 100 * class_correct[i] / class_total[i]))